#include "socket_proxy.h"

#include <core/QbLogger.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <sstream>

#include "base64.h"
#include "socket_inc.h"
#include "socket_wrapper.h"

#ifdef _MSC_VER
#include <WinInet.h>
#include <tchar.h>
#pragma comment(lib, "wininet.lib")
#endif

namespace qb {
ProxyInfo CSocketProxy::g_info;

CSocketProxy::CSocketProxy(void) {}

CSocketProxy::~CSocketProxy(void) {}

int32_t CSocketProxy::Test(ProxyType proxy_type, const char* dst_host,
                           uint16_t dst_port, const char* proxy_host,
                           uint16_t proxy_port, const char* username,
                           const char* password, const char* domain) {
  CSocketProxy proxy;
#ifdef _MSC_VER
  uint32_t tSocket = INVALID_SOCKET;
#else
  int32_t tSocket = INVALID_SOCKET;
#endif
  hostent* pHostent = NULL;
  sockaddr_in clientAddr;

  proxy.SetProxy(proxy_type, proxy_host, proxy_port, username, password,
                 domain);
  pHostent = gethostbyname(proxy.GetIPAddress());
  if (!pHostent) {
    log_error("gethostbyname(%s) falied, getlasterror(%d)",
              proxy.GetIPAddress(), CSocketWrapper::GetSockError());
    return -1;
  }
  clientAddr.sin_family = AF_INET;
  clientAddr.sin_port = htons(proxy.GetPort());
  clientAddr.sin_addr = *((struct in_addr*)pHostent->h_addr);

  tSocket = (uint32_t)::socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
  if (tSocket == SOCKET_ERROR) {
    return -1;
  }
  if (connect(tSocket, (sockaddr*)&clientAddr, sizeof(clientAddr)) ==
      SOCKET_ERROR) {
#ifdef _MSC_VER
    closesocket(tSocket);
#else
    close(tSocket);
#endif
    return -1;
  }
  if (proxy.Connect(tSocket, dst_host, dst_port, false) == -1) {
#ifdef _MSC_VER
    closesocket(tSocket);
#else
    close(tSocket);
#endif
    return -1;
  }

  return 0;
}

void CSocketProxy::SetGlobalProxy(ProxyType proxy_type, const char* proxy_host,
                                  uint16_t proxy_port, const char* username,
                                  const char* password, const char* domain) {
  SetProxyInfo(g_info, proxy_type, proxy_host, proxy_port, username, password,
               domain);
}

void CSocketProxy::SetProxy(ProxyType proxy_type, const char* proxy_host,
                            uint16_t proxy_port, const char* username,
                            const char* password, const char* domain) {
  SetProxyInfo(m_info, proxy_type, proxy_host, proxy_port, username, password,
               domain);
}

void CSocketProxy::SetProxyInfo(ProxyInfo& info, ProxyType proxy_type,
                                const char* proxy_host, uint16_t proxy_port,
                                const char* username, const char* password,
                                const char* domain) {
  info.bUse = true;
  info.nProxyPort = proxy_port;
  info.nProxyType = proxy_type;
  info.strProxyHost = proxy_host == NULL ? "" : proxy_host;
  info.strUsername = username == NULL ? "" : username;
  info.strPassword = password == NULL ? "" : password;
  info.strDomain = domain == NULL ? "" : domain;
#ifdef _MSC_VER
  if (proxy_type == kProxyBrowser) {
    std::wstring strIP;
    std::wstring strUsername;
    std::wstring strPassword;
    uint16_t port = 0;
    if (GetBrowserInfo(strIP, port, strUsername, strPassword) == 0) {
      int len8 = ::WideCharToMultiByte(
          CP_ACP, 0, strIP.c_str(), (int)strIP.length(), NULL, 0, NULL, NULL);
      len8 = len8 > 1024 ? 1024 : len8;
      char ns[1024 + 4] = {0};
      ::WideCharToMultiByte(CP_ACP, 0, strIP.c_str(), (int)strIP.length(), ns,
                            len8, NULL, NULL);
      info.strProxyHost = ns;
      info.nProxyPort = port;
    } else {
      info.nProxyType = kProxyNone;
    }
  }
#endif
}

ProxyType CSocketProxy::GetProxyType() {
  return (ProxyType)GetInfo().nProxyType;
}

const ProxyInfo& CSocketProxy::GetInfo() {
  if (m_info.bUse) {
    return m_info;
  }

  return g_info;
}

uint16_t CSocketProxy::GetPort() { return GetInfo().nProxyPort; }

const char* CSocketProxy::GetIPAddress() {
  return GetInfo().strProxyHost.c_str();
}

#ifdef _MSC_VER
bool CSocketProxy::IsIEVerAbove5() {
  HKEY hCU;

  if (RegOpenKeyEx(HKEY_LOCAL_MACHINE,
                   _T("Software\\Microsoft\\Internet Explorer"), 0,
                   KEY_QUERY_VALUE, &hCU) == ERROR_SUCCESS) {
    char szIEx[100] = {0};
    DWORD ulSize = 100;

    if (RegQueryValueEx(hCU, _T("Version"), NULL, NULL, (LPBYTE)szIEx,
                        &ulSize) == ERROR_SUCCESS) {
      std::wstring ver = (LPCTSTR)szIEx;
      std::wstring majVer;
      std::size_t nPos = ver.find(_T('.'));
      if (nPos >= 0) {
        majVer = ver.substr(0, ver.find('.'));
        if (_ttoi(majVer.c_str()) >= 5) {
          return true;
        }
      }
    }
  }

  return false;
}
#endif

int32_t CSocketProxy::GetBrowserInfo(std::wstring& strIP, uint16_t& port,
                                     std::wstring& strUsername,
                                     std::wstring& strPassword) {
#ifdef _MSC_VER
  std::wstring strProxyLists = _T("");
  bool bUseIE5 = IsIEVerAbove5();
  if (bUseIE5) {
    INTERNET_PER_CONN_OPTION_LIST List;
    INTERNET_PER_CONN_OPTION Option[5];
    unsigned long nSize = sizeof(INTERNET_PER_CONN_OPTION_LIST);

    Option[0].dwOption = INTERNET_PER_CONN_AUTOCONFIG_URL;
    Option[1].dwOption = INTERNET_PER_CONN_AUTODISCOVERY_FLAGS;
    Option[2].dwOption = INTERNET_PER_CONN_FLAGS;
    Option[3].dwOption = INTERNET_PER_CONN_PROXY_BYPASS;
    Option[4].dwOption = INTERNET_PER_CONN_PROXY_SERVER;

    List.dwSize = sizeof(INTERNET_PER_CONN_OPTION_LIST);
    List.pszConnection = NULL;
    List.dwOptionCount = 5;
    List.dwOptionError = 0;
    List.pOptions = Option;

    if (::InternetQueryOption(NULL, INTERNET_OPTION_PER_CONNECTION_OPTION,
                              &List, &nSize)) {
      if ((Option[2].Value.dwValue & PROXY_TYPE_PROXY) == PROXY_TYPE_PROXY) {
        strProxyLists = Option[4].Value.pszValue;
      }

      if (Option[0].Value.pszValue != NULL) {
        GlobalFree(Option[0].Value.pszValue);
      }
      if (Option[3].Value.pszValue != NULL) {
        GlobalFree(Option[3].Value.pszValue);
      }
      if (Option[4].Value.pszValue != NULL) {
        GlobalFree(Option[4].Value.pszValue);
      }
    }
  }

  if (strProxyLists.empty()) {
    bool bEnableDial = false;
    HKEY hk;
    LONG lRet = RegOpenKeyEx(
        HKEY_CURRENT_USER,
        _T("Software\\Microsoft\\Windows\\CurrentVersion\\Internet Settings"),
        NULL, KEY_READ | KEY_QUERY_VALUE, &hk);
    if (lRet == ERROR_SUCCESS && NULL != hk) {
      uint32_t dwDialEnable = 0;

      uint32_t dwType = REG_DWORD;
      DWORD dwSize = sizeof(dwDialEnable);
      lRet = RegQueryValueEx(hk, _T("EnableAutodial"), NULL, NULL,
                             (LPBYTE)&dwDialEnable, &dwSize);

      if (lRet == ERROR_SUCCESS && dwDialEnable != 0) {
        bEnableDial = true;
      }
      RegCloseKey(hk);
    }
  } else {
    while (!strProxyLists.empty()) {
      std::size_t nPos = strProxyLists.find(';', 0);
      std::wstring strItem;
      if (nPos == -1) {
        strItem = strProxyLists;
        strProxyLists = L"";
      } else {
        strItem = strProxyLists.substr(0, nPos);
        strProxyLists = strProxyLists.substr(nPos + 1, strProxyLists.length());
      }

      std::size_t nPos2 = strItem.find('=', 0);

      std::wstring strType;
      if (nPos2 != -1) {
        strType = strItem.substr(0, nPos2);
        strItem = strItem.substr(nPos2 + 1, strItem.length());
      }

      // 读出第一个代理信息的数据
      bool bFoundProxy = false;

      if (strType == _T("http") || strType.empty()) {
        // 使用http代理
        std::size_t nPos3 = strItem.find(':');
        strIP = strItem.substr(0, nPos3);
        strItem = strItem.substr(nPos3 + 1, strItem.length());
        port = _wtoi(strItem.c_str());

        DWORD dwLen = 64;
        WCHAR szUserName[64] = {0};
        WCHAR szPassword[64] = {0};
        ::InternetQueryOption(NULL, INTERNET_OPTION_PROXY_USERNAME, szUserName,
                              &dwLen);

        dwLen = 64;
        ::InternetQueryOption(NULL, INTERNET_OPTION_PROXY_PASSWORD, szPassword,
                              &dwLen);

        strUsername = szUserName;
        strPassword = szPassword;

        bFoundProxy = true;
      } else if (strType == _T("socks")) {
        // 使用socks代理
        std::size_t nPos3 = strItem.find(':');
        strIP = strItem.substr(0, nPos3);
        strItem = strItem.substr(nPos3 + 1, strItem.length());
        port = _wtoi(strItem.c_str());

        bFoundProxy = true;
      }

      if (bFoundProxy) {
        return 0;
      }
    }
  }
#endif
  return -1;
}

int32_t CSocketProxy::Select(uint32_t tSocket) {
  fd_set fdRead;
  FD_ZERO(&fdRead);
  FD_SET(tSocket, &fdRead);
  timeval tvSelect_Time_Out;
  tvSelect_Time_Out.tv_sec = 3;
  tvSelect_Time_Out.tv_usec = 0;
  int32_t nRet = select(tSocket + 1, &fdRead, NULL, NULL, &tvSelect_Time_Out);
  if (nRet <= 0) {
    return -1;
  }
  return 0;
}

int32_t CSocketProxy::UdpAssociate(uint32_t tSocket, const char* ipaddress,
                                   uint16_t port, std::string& proxy_reply,
                                   uint16_t& proxy_reply_port) {
  // 发送UDP 穿透信令。
  /**
  +----+-----+-------+------+----------+----------+
  |VER | CMD |  RSV  | ATYP | DST.ADDR | DST.PORT |
  +----+-----+-------+------+----------+----------+
  | 1  |  1  | X'00' |  1   | Variable |    2     |
  +----+-----+-------+------+----------+----------+
  Where:
  o  VER    protocol version: X'05'
  o  CMD
  o  CONNECT X'01'
  o  BIND X'02'
  o  UDP ASSOCIATE X'03'
  o  RSV    RESERVED
  o  ATYP   address type of following address
  o  IP V4 address: X'01'
  o  DOMAINNAME: X'03'
  o  IP V6 address: X'04'
  o  DST.ADDR       desired destination address
  o  DST.PORT desired destination port in network octet order
  */
  udpassociatereq associate;
  associate.Ver = kProxySocksVer5;
  associate.Cmd = 0x03;
  associate.Rsv = 0;
  associate.Atyp = 0x01;
  associate.IPAddr = inet_addr(ipaddress);
  associate.Port = htons(port);

  if (SOCKET_ERROR == send(tSocket, (char*)&associate, sizeof(associate), 0)) {
    return -1;
  }

  // 接收回应。
  /**
  +----+-----+-------+------+----------+----------+
  |VER | REP |  RSV  | ATYP | BND.ADDR | BND.PORT |
  +----+-----+-------+------+----------+----------+
  | 1  |  1  | X'00' |  1   | Variable |    2     |
  +----+-----+-------+------+----------+----------+
  Where:
  o  VER    protocol version: X'05'
  o  REP    Reply field:
  o  X'00' succeeded
  o  X'01' general SOCKS server failure
  o  X'02' connection not allowed by ruleset
  o  X'03' Network unreachable
  o  X'04' Host unreachable
  o  X'05' Connection refused
  o  X'06' TTL expired
  o  X'07' Command not supported
  o  X'08' Address type not supported
  o  X'09' to X'FF' unassigned
  o  RSV    RESERVED
  o  ATYP   address type of following address
  */

  if (Select(tSocket) == -1) {
    return -1;
  }
  char buff[512] = {0};
  struct udpassociateans* proxyans = (udpassociateans*)buff;
  int32_t len = ::recv(tSocket, buff, sizeof(buff), 0);
  if (len <= 0 || kProxySocksVer5 != proxyans->Ver || 0 != proxyans->Rep ||
      proxyans->Atyp != 1) {
    return -1;
  }
  proxy_reply = inet_ntoa(*(in_addr*)(&proxyans->IPAddr));
  proxy_reply_port = ntohs(*(short*)(&proxyans->Port));

  return 0;
}

udpheader CSocketProxy::GetUdpHeader(const char* pIPAddress, uint16_t port) {
  // 发送UDP数据要附加头。
  /**
  +----+------+------+----------+----------+----------+
  |RSV | FRAG | ATYP | DST.ADDR | DST.PORT |   DATA   |
  +----+------+------+----------+----------+----------+
  | 2  |  1   |  1   | Variable |    2     | Variable |
  +----+------+------+----------+----------+----------+
  The fields in the UDP request header are:
  o  RSV  Reserved X'0000'
  o  FRAG    Current fragment number
  o  ATYP    address type of following addresses:
  o  IP V4 address: X'01'
  o  DOMAINNAME: X'03'
  o  IP V6 address: X'04'
  o  DST.ADDR       desired destination address
  o  DST.PORT       desired destination port
  o  DATA     user data
  */

  udpheader header;
  header.Rsv = 0;
  header.Frag = 0;
  header.Atyp = 0x01;
  header.IPAddr = inet_addr(pIPAddress);
  header.Port = htons(port);

  return header;
}

int32_t CSocketProxy::Connect(uint32_t tSocket, const char* host, uint16_t port,
                              bool use_udp) {
  // 验证socks5服务器
  if (GetInfo().nProxyType == kProxySocksVer5 ||
      GetInfo().nProxyType == kProxySocksVer4) {
    std::string remote_addr;
    int32_t ret =
        ConnectSocks5(tSocket, host, port, GetInfo().strProxyHost.c_str(),
                      GetInfo().nProxyPort, GetInfo().strUsername.c_str(),
                      GetInfo().strPassword.c_str(), remote_addr,
                      m_remoteInfo.uRemoteProxyPort, use_udp);
    m_remoteInfo.strRemoteAddr = remote_addr;

    return ret;
  } else if (GetInfo().nProxyType == kProxyHttp ||
             GetInfo().nProxyType == kProxyBrowser) {
    return ConnectHTTP(tSocket, host, port, GetInfo().strProxyHost.c_str(),
                       GetInfo().nProxyPort, GetInfo().strUsername.c_str(),
                       GetInfo().strPassword.c_str());
  } else {
    return -1;
  }
}

int32_t CSocketProxy::ConnectHTTP(uint32_t tSocket, const char* host,
                                  uint16_t port, const char* proxy_host,
                                  uint16_t proxy_port, const char* username,
                                  const char* password) {
  std::string str;
  std::stringstream ss_stream;
  ss_stream << "CONNECT ";
  ss_stream << host;
  ss_stream << ":";
  ss_stream << port;
  ss_stream << " HTTP/1.1\r\nHost: ";
  ss_stream << host;
  ss_stream << ":";
  ss_stream << port;
  ss_stream << "\r\n";
  if (username && std::string(username).length() > 0) {
    std::string strUserPass;
    strUserPass.append(username);
    strUserPass.append(":");
    strUserPass.append(password);
    char strBase64[256] = {0};
    Base64_Encode(strBase64, strUserPass.c_str(), strUserPass.length());

    ss_stream << "Proxy-Authorization: Basic ";
    ss_stream << strBase64;
    ss_stream << "\r\n";
  }
  str = ss_stream.str();
  str.append("Accept: */*\r\n");
  str.append("Content-Type: text/html\r\n");
  str.append("Proxy-Connection: Keep-Alive\r\n");
  str.append("Content-length: 0\r\n");
  str.append("\r\n");
  send(tSocket, str.c_str(), (int)str.length(), 0);

  if (Select(tSocket) == -1) {
    return -1;
  }
  char buff[512] = {0};
  int32_t len = recv(tSocket, buff, sizeof(buff), 0);
  if (len <= 0) {
    return -1;
  }
  if (strstr(buff, "HTTP/1.1 200") != NULL ||
      strstr(buff, "HTTP/1.0 200") != NULL) {
    return 0;
  }

  return -1;
}

int32_t CSocketProxy::ConnectSocks5(uint32_t tSocket, const char* host,
                                    uint16_t port, const char* proxy_host,
                                    uint16_t proxy_port, const char* username,
                                    const char* password,
                                    std::string remote_addr,
                                    uint16_t& remote_proxy_port, bool use_udp) {
  hostent* pHostent = gethostbyname(host);
  if (pHostent == NULL) {
    log_error("gethostbyname(%s) falied, getlasterror(%d)", host,
              CSocketWrapper::GetSockError());
    return -1;
  }

  // Step 1: 连接代理服务器成功后，马上开始和代理协商，
  // 协商报文如下,询问服务器，版本5，方法个数,是需要验证(0x02)还是不需要验证(0x00)
  // +--------+-------------------+---------------+
  // | VER    | Number of METHODS | METHODS       |
  // +--------+-------------------+---------------+
  // | 0x05   | 0x02 (有两个方法) | 0x00 | 0x02   |
  // +--------+-------------------+---------------+
  struct socks5req1 m_proxyreq1;
  m_proxyreq1.Ver = kProxySocksVer5;
  m_proxyreq1.nMethods = 2;
  m_proxyreq1.Methods[0] = 0;
  m_proxyreq1.Methods[1] = 2;
  send(tSocket, (char*)&m_proxyreq1, sizeof(m_proxyreq1), 0);

  // Setp 2: 代理服务器将返回两个字节的协商结果，接收协商结果。
  if (Select(tSocket) == -1) {
    return -1;
  }
  char buff[512] = {0};
  struct socks5ans1* proxyans1 = (socks5ans1*)buff;
  recv(tSocket, buff, sizeof(buff), 0);
  if (proxyans1->Ver != 5 ||
      (proxyans1->Method != 0 && proxyans1->Method != 2)) {
    return -1;
  }

  // Step 3: 根据协商结果判断是否需要验证用户，如果是0x02，则需要提供验证，
  // 验证部分参考RFC1929
  if (proxyans1->Method == 2 && strlen(username) > 0) {
    uint8_t nUserLen = (uint8_t)strlen(username);
    uint8_t nPassLen = (uint8_t)strlen(password);

    // Step 3-1:RFC1929，初始协商 参考:
    // http://cnlinzhimao.blog.163.com/blog/static/2352787220080763150619/
    // +-----------------------------------------------+
    // | VER    | ULEN  | UNAME     | PLEN  | PASSWD   |
    // +-----------------------------------------------+
    // | 1      | 1     | 1 to 255  | 1     | 1 to 255 |
    // +-----------------------------------------------+
    struct authreq m_authreq;
    memset(&m_authreq, 0, sizeof(m_authreq));
    m_authreq.Ver = 0x01;
    m_authreq.Ulen = nUserLen;
#ifdef _MSC_VER
    strncpy_s(m_authreq.NamePass,
              sizeof(m_authreq.NamePass) / sizeof(m_authreq.NamePass[0]),
              username, _TRUNCATE);
#else
    strcpy(m_authreq.NamePass, username);
#endif
    memcpy(m_authreq.NamePass + nUserLen, &nPassLen, 1);
#ifdef _MSC_VER
    strncpy_s(m_authreq.NamePass + nUserLen + 1,
              sizeof(m_authreq.NamePass) / sizeof(m_authreq.NamePass[0]) -
                  nUserLen - 1,
              password, _TRUNCATE);
#else
    strcpy(m_authreq.NamePass + nUserLen + 1, password);
#endif
    int32_t len = 3 + nUserLen + nPassLen;
    send(tSocket, (char*)&m_authreq, len, 0);

    // Step 3-2:账号密码验证结果
    struct authans* m_authans = (struct authans*)buff;
    memset(buff, 0, sizeof(buff));

    if (Select(tSocket) == -1) {
      return -1;
    }
    recv(tSocket, buff, sizeof(buff), 0);
    if (m_authans->Status != 0x00) {
      return -1;
    }
  }
  if (use_udp)  // udp不需要继续连接了，否则无法使用udp
  {
    return 0;
  }

  // Step 4: 协商完成，开始发送连接远程服务器请求,请求报文格式如下:
  // +----+-----+-------+------+----------+----------+
  // |VER | CMD |　RSV　| ATYP | DST.ADDR | DST.PORT |
  // +----+-----+-------+------+----------+----------+
  // | 1　| 　1 | 0x00  | 　1　| Variable |　　 2	   |
  // +----+-----+-------+------+----------+----------+
  // CMD==0x01 表示连接,
  // ATYP==0x01表示采用IPV4格式地址，DST.ADDR是远程服务器地址，DST.PORT是远程服务器端口
  // 如果需要接受外来连接，则需要在连接完成之后，发送CMD==0x02绑定请求，代理将为此请求绑定一个套接字接受外部连接
  struct socks5req2 m_proxyreq2;
  m_proxyreq2.Ver = 5;
  m_proxyreq2.Cmd = 1;
  m_proxyreq2.Rsv = 0;
  m_proxyreq2.Atyp = 1;
  m_proxyreq2.IPAddr = *(uint32_t*)pHostent->h_addr_list[0];
  m_proxyreq2.Port = ntohs(port);
  send(tSocket, (char*)&m_proxyreq2, 10, 0);

  // Step 5: 接收对请求的响应，响应包格式如下：
  // +----+-----+-------+------+----------+----------+
  // |VER | REP |　RSV　| ATYP | BND.ADDR | BND.PORT |
  // +----+-----+-------+------+----------+----------+
  // | 1　|　1　| 0x00  |　1 　| Variable | 　　2 　 |
  // +----+-----+-------+------+----------+----------+
  // VER 必须是0x05, REP==0x00表示成功，ATYP==0x01表示地址是IPV4地址，
  // BND.ADDR 是代理为连接远程服务器绑定的地址，BND.PORT是这个套接字的端口
  struct socks5ans2* proxyans2;
  proxyans2 = (struct socks5ans2*)buff;
  memset(buff, 0, sizeof(buff));

  if (Select(tSocket) == -1) {
    return -1;
  }
  int32_t nRet = recv(tSocket, buff, sizeof(buff), 0);
  if (nRet == -1) {
    return -1;
  }
  if (proxyans2->Ver != 5 || proxyans2->Rep != 0) {
    return -1;
  }
#ifdef _MSC_VER
  if (nRet >= 10) {
    // 得到代理绑定地址
    struct sockaddr_in m_SockBindAddr4;
    m_SockBindAddr4.sin_addr.S_un.S_addr =
        *(unsigned long*)&proxyans2->other[0];

    m_SockBindAddr4.sin_port =
        *(unsigned short*)(&proxyans2->other[0] + sizeof(long));
    unsigned char* addr = (unsigned char*)&m_SockBindAddr4.sin_addr.s_addr;

    std::stringstream ss_stream;
    ss_stream << addr[0];
    ss_stream << ".";
    ss_stream << addr[1];
    ss_stream << ".";
    ss_stream << addr[2];
    ss_stream << ".";
    ss_stream << addr[3];
    remote_addr = ss_stream.str();
    remote_proxy_port = ntohs(m_SockBindAddr4.sin_port);
  }
#endif
  return 0;
}

}  // namespace qb
